package com.eonmind.ai.core.chat.advisors.memory;

import com.eonmind.ai.converts.AiChatConvert;
import com.eonmind.ai.service.impl.AiConversationService;
import com.eonmind.ai.service.impl.AiMessageService;
import com.eonmind.ai.service.impl.AiModelService;
import com.eonmind.common.contract.application.redis.service.RedisService;
import com.eonmind.common.provider.entity.ai.AiConversation;
import com.eonmind.common.provider.entity.ai.AiConversationMessage;
import com.eonmind.common.provider.entity.ai.AiModel;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.Message;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Lazy;
import org.springframework.stereotype.Component;

import java.util.List;

import static com.eonmind.common.contract.constants.ai.AiConstant.AI_CONVERSATION_MODEL_REDIS_KEY;

@Slf4j
@Component
public class MongoChatMemory implements ChatMemory {

    @Resource
    private AiMessageService aiMessageService;

    @Resource
    private RedisService redisService;

    @Resource
    private AiModelService aiModelService;

    @Lazy
    @Resource
    private AiConversationService aiConversationService;

    @Override
    public void add(String conversationId, List<Message> messages) {
        List<AiConversationMessage> chatMessageList = messages.stream()
                .map(t -> {
                    Long conversationIdLong = Long.valueOf(conversationId);
                    // 查询对应的回话信息
                    AiConversation conversation = aiConversationService.selectConversationById(conversationIdLong);

                    AiConversationMessage aiConversationMessage = AiChatConvert.toAiConversationMessage(conversation, t);
                    Object modelId = redisService.get(AI_CONVERSATION_MODEL_REDIS_KEY + conversationId);
                    if (modelId != null) {
                        try {
                            Long id = Long.valueOf(modelId.toString());
                            AiModel aiModel = aiModelService.selectById(id);
                            aiConversationMessage.setModelId(aiModel.getId());
                            aiConversationMessage.setModelName(aiModel.getModelName());
                        } catch (NumberFormatException e) {
                            log.error("ModelId is not a number: {}", modelId);
                        }
                    }
                    return aiConversationMessage;
                }).toList();
        aiMessageService.saveBatch(chatMessageList);
    }

    @Override
    public List<Message> get(String conversationId, int lastN) {
        List<AiConversationMessage> chatMessageList = aiMessageService.selectByConversationIdAndLimit(Long.valueOf(conversationId), lastN);
        return chatMessageList.stream()
                .map(AiChatConvert::toMessageByType).toList();
    }

    @Override
    public void clear(String conversationId) {
        aiMessageService.deleteByConversationId(Long.valueOf(conversationId));
    }

    @Bean(name = "mongoChatMemoryAdvisor")
    public MessageChatMemoryAdvisor mongoMessageChatMemoryAdvisor(
            ChatMemory mongoChatMemory
    ) {

        return new MessageChatMemoryAdvisor(mongoChatMemory);
    }
}
